import argparse
import itertools
import os
import numpy as np
import torch
import triton
import triton.language as tl
import triton.language.extra.libdevice as tldevice

FLA_SINGLE_THREAD = os.getenv("TRITON_PRECISION_NPU_SINGLE_THREAD", '0') == '1'

def single_wrapper(grid, grid_size, kernel):   
    if not FLA_SINGLE_THREAD:
        kwargs = {'USE_PIDS': False}
        for dim in range(grid_size):
            kwargs[f"PID{dim}"] = 0
        kernel(grid, kwargs)
        return
    assert isinstance(grid, tuple)
    
    grid_ranges = []
    for dim in grid:
        grid_ranges.append(range(dim))
    single_grid = (1,) * grid_size

    for combination in itertools.product(*grid_ranges):
        kwargs = {'USE_PIDS': True}
        for idx, dim in enumerate(combination):
            kwargs[f"PID{idx}"] = dim
        kernel(single_grid, kwargs)

if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':
    exp = tldevice.fast_expf
    exp2 = tldevice.exp2
    log = tldevice.fast_logf
    log2 = tldevice.fast_log2f
else:
    exp = tl.exp
    exp2 = tl.math.exp2
    log = tl.log
    log2 = tl.log2

@triton.heuristics({
    'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None,
    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
})
@triton.jit(do_not_specialize=['PID0', 'PID1'])
def causal_conv1d_states_fwd_kernel(
    x,
    initial_state,
    final_state,
    cu_seqlens,
    T,
    D,
    W,
    BD: tl.constexpr,
    BW: tl.constexpr,
    USE_INITIAL_STATE: tl.constexpr,
    IS_VARLEN: tl.constexpr,
    USE_PIDS: tl.constexpr,
    PID0,
    PID1
):
    if USE_PIDS:
        i_d, i_n = PID0, PID1
    else:
        i_d, i_n = tl.program_id(0), tl.program_id(1)
    if IS_VARLEN:
        bos, eos = tl.load(cu_seqlens + i_n), tl.load(cu_seqlens + i_n + 1)
        T = eos - bos
    else:
        bos, eos = i_n * T, i_n * T + T

    o_t = eos - BW + tl.arange(0, BW)
    o_d = i_d * BD + tl.arange(0, BD)
    o_w = W - BW + tl.arange(0, BW)
    m_t = (o_t >= tl.maximum(bos, eos - W))
    m_d = o_d < D
    m_w = (o_w >= 0) & (o_w < W)

    b_x = tl.load(x + o_t * D + o_d[:, None], mask=(m_t & m_d[:, None]), other=0)
    if USE_INITIAL_STATE:
        if T < BW:
            o_c = W - (BW - T) + tl.arange(0, BW)
            m_c = (o_c >= 0) & (o_c < W)
            b_cache = tl.load(initial_state + i_n * D*W + o_d[:, None] * W + o_c, mask=m_d[:, None] & m_c, other=0)
            b_x += b_cache

    tl.store(final_state + i_n * D*W + o_d[:, None] * W + o_w, b_x, mask=m_d[:, None] & m_w)

def test_causal_conv1d_states_fwd_kernel(output_file):
    # 设置随机种子以确保可重复性
    torch.manual_seed(42)

    # 定义参数
    B = 2        # 批量大小
    T = 64       # 序列长度
    D = 32       # 特征维度
    W = 16       # 卷积窗口大小
    BD = 8       # block大小 for D
    BW = 8       # block大小 for W

    # 生成随机输入张量
    device = 'npu'  # 使用NPU设备
    dtype = torch.float32  # 使用float32以匹配内核的内部类型

    # 输入张量
    x = torch.randn(B, T, D, dtype=dtype, device=device)
    initial_state = torch.randn(B, D, W, dtype=dtype, device=device)
    final_state = torch.randn(B, D, W, dtype=dtype, device=device)
    cu_seqlens = torch.randint(low=0, high=T, size=(B + 1,), dtype=torch.int32, device=device)
    cu_seqlens[0] = 0
    cu_seqlens[-1] = T

    # 计算网格大小
    num_blocks_d = triton.cdiv(D, BD)
    num_blocks_n = B
    grid = (num_blocks_d, num_blocks_n)

    # 启用功能标志
    USE_INITIAL_STATE = True
    IS_VARLEN = True

    # 调用内核函数
    def kernel(cur_grid, pid_args):
        causal_conv1d_states_fwd_kernel[cur_grid](
            x, initial_state, final_state, cu_seqlens,
            T=T, D=D, W=W, BD=BD, BW=BW,
            USE_INITIAL_STATE=USE_INITIAL_STATE,
            IS_VARLEN=IS_VARLEN,
            **pid_args
        )
    single_wrapper(grid, 2, kernel)

    y_numpy = final_state.cpu().detach().numpy()
    np.savetxt(output_file, y_numpy.reshape(-1, y_numpy.shape[-1]))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Test Causal Conv1D Update Kernel')
    parser.add_argument('--output', type=str, default='default_output.txt', 
                        help='Output file name (default: default_output.txt)')
    args = parser.parse_args()
    test_causal_conv1d_states_fwd_kernel(args.output)
